from numpy import ndarray
from scipy import ndimage
from scipy.ndimage import filters
import numpy as np
import logging

logger = logging.getLogger('pyastm.spectrogram')

class Spectrogram:
    def __init__(self, freq_snapshots, freq_band, time_duration):
        self.freq_snapshots = freq_snapshots
        self.freq_band = freq_band
        self.time_duration = time_duration
        logger.debug('new Spectrogram instance: freq_band: %s, time_duration: %s' % (str(freq_band), str(time_duration)))
        
    def resize(self, size):
        return resize(self, size)
    
    def change_freq_band(self, freq_band_new, keep_original_size=True):
        return change_freq_band(self, freq_band_new, keep_original_size)
            
    def resize_and_change_freq_band(self, size, freq_band_new=None):
        return resize_and_change_freq_band(self, size, freq_band_new)
    
    def get_fbpp(self):
        '''Returns frequency band per pixel'''
        return float(self.freq_band) / self.freq_snapshots.shape[0]
    
    def get_fsps(self):
        '''Returns frequency snapshots per second'''
        return float(self.freq_snapshots.shape[1]) / self.time_duration
    
    def get_size(self):
        return self.freq_snapshots.shape
    

def resize(spectrogram, size):
    freq_snapshots_resized = _resize_array(spectrogram.freq_snapshots, size)
    return Spectrogram(freq_snapshots_resized, spectrogram.freq_band, spectrogram.time_duration)
    
def resize_and_change_freq_band(spectrogram, size, freq_band_new):
    freq_band = spectrogram.freq_band
    freq_snapshots = spectrogram.freq_snapshots
    
    # check the argument
    if freq_band_new == None:
        raise ValueError("freq_band_new mustn not be None")
    
    if freq_band_new > freq_band:
        # expand the frequency band with zeros
        freq_snapshots_new = _expand_freq_band(spectrogram, freq_band_new)
    elif freq_band_new < freq_band:
        # cut a sub-band of frequencies from the original
        subarray_width = np.round((freq_band_new / freq_band) * freq_snapshots.shape[0])
        freq_snapshots_new = freq_snapshots[-subarray_width::, :]
    else:
        # frequency band stays unchanged
        freq_snapshots_new = freq_snapshots
    freq_snapshots_resized = _resize_array(freq_snapshots_new, size)
    return Spectrogram(freq_snapshots_resized, freq_band_new, spectrogram.time_duration)

def change_freq_band(spectrogram, freq_band_new, keep_original_size=True):
    if spectrogram.freq_band == freq_band_new:
        return spectrogram
    
    # spectrogram data
    freq_band = spectrogram.freq_band
    freq_snapshots = spectrogram.freq_snapshots
    time_duration = spectrogram.time_duration
    
    if freq_band_new > freq_band:
        # expand the frequency band with zeros
        complemented_array = _expand_freq_band(spectrogram, freq_band_new)
        if ~keep_original_size:
            return Spectrogram(complemented_array, freq_band_new, time_duration)
        else:
            freq_snapshot_new = complemented_array
    else:
        # cut a sub=band from the original spectrogram
        subarray_width = np.round((freq_band_new / freq_band) * freq_snapshots.shape[0])
        subarray = freq_snapshots[-subarray_width::, :]
        if ~keep_original_size:
            return Spectrogram(subarray, freq_band_new, time_duration)
        else:
            freq_snapshot_new = subarray
            
    freq_snapshots_resized = _resize_array(freq_snapshot_new, (freq_snapshots.shape))
    return Spectrogram(freq_snapshots_resized, freq_band_new, time_duration)
        
    
def create_spectrogram_for_signal(signal, signal_sampling_freq, fftlen, freq_snapshots_per_sec):
    freq_snapshots, time_duration = stft_spaced(signal, fftlen, signal_sampling_freq, freq_snapshots_per_sec)
    freqs = np.fft.fftfreq(fftlen, 1.0 / signal_sampling_freq)
    return Spectrogram(freq_snapshots, np.max(np.abs(freqs)), time_duration)
            
def stft_full(signal, fftlen):
    frequency_snapshots = np.zeros((fftlen / 2, len(signal)), dtype=np.float)
    for i in range(len(signal) - fftlen + 1):
        signal_snippet = signal[i:i + fftlen - 1]
        frequency_snapshots[::-1, i] = abs(np.fft.rfft(signal_snippet))
    return frequency_snapshots
            
def stft_spaced(signal, fftlen, signal_sampling_freq, freq_snapshots_per_sec):
    # step between two signal snippets
    step = float(signal_sampling_freq) / freq_snapshots_per_sec
    
    # each index is a starting point for one signal snippet
    indices = ndarray.round(np.arange(0, len(signal) - fftlen + 1, step))
    
    # prepare the output array
    frequency_snapshots = np.zeros((fftlen / 2, len(indices)), dtype=np.float)
    
    # compute the duration of the spectrogram
    time_duration = step / float(signal_sampling_freq) * len(indices)
    logger.debug('resulting spectrogram time duration: %s seconds' % str(time_duration))
    
    # compute the fft from every signal snippet
    for i in range(len(indices)):
        index = indices[i]
        signal_snippet = signal[index:index + fftlen - 1]
        frequency_snapshots[::-1, i] = np.abs(np.fft.rfft(signal_snippet))
    return frequency_snapshots, time_duration

def _resize_array(array, size, prefilter=True):
    zoom_x = float(size[1]) / array.shape[1]
    zoom_y = float(size[0]) / array.shape[0]
    
    # prefilter with gaussian filter along each shrinked axis
    if prefilter:
        if zoom_y < 0:
            sigma_y = 0.5 / zoom_y
            array = filters.gaussian_filter1d(array, sigma_y, 0, 0, None, 'constant', 0)
        if zoom_x < 0:
            sigma_x = 0.5 * zoom_x
            array = filters.gaussian_filter1d(array, sigma_x, -1, 0, None, 'constant', 0)
    
    return ndimage.zoom(array, [zoom_y, zoom_x])

def _expand_freq_band(spectrogram, freq_band_new):
    new_rows_count = np.round((freq_band_new - spectrogram.freq_band) / spectrogram.get_fbpp())
    new_rows = np.zeros((new_rows_count, spectrogram.freq_snapshots.shape[1]))
    return np.concatenate([new_rows, spectrogram.freq_snapshots])